package com.zhang.spark_2.com.zhang.core.acc

import org.apache.spark.rdd.RDD
import org.apache.spark.util.AccumulatorV2
import org.apache.spark.{SparkConf, SparkContext}

import scala.collection.mutable

/**
 * @title:
 * @author: zhang
 * @date: 2022/2/16 18:36 
 */
object Spark02_Acc {

  def main(args: Array[String]): Unit = {

    //TODO 获取环境
    val conf = new SparkConf().setMaster("local").setAppName("WordCount")
    val sc = new SparkContext(conf);

    val list: RDD[String] = sc.makeRDD(List("a", "a", "b", "b", "b"), 2)

    //注册累加器
    val acc = new MyAcc()
    sc.register(acc, "myacc")
    list.foreach(
      word => {
        acc.add(word)
      }
    )
    println(acc.value)


    sc.stop()

  }
}

class MyAcc extends AccumulatorV2[String, mutable.Map[String, Int]] {
  val map1 = mutable.Map[String, Int]()

  override def isZero: Boolean = {
    map1.isEmpty
  }

  override def copy(): AccumulatorV2[String, mutable.Map[String, Int]] = {
    new MyAcc()
  }

  override def reset(): Unit = {
    map1.clear()
  }

  override def add(word: String): Unit = {
    val oldCnt: Int = map1.getOrElse(word, 0)
    map1.update(word, oldCnt + 1)
  }

/*  override def merge(other: AccumulatorV2[String, mutable.Map[String, Int]]): Unit = {
    val map2: mutable.Map[String, Int] = other.value
    map2.foreach {
      case (word, cnt) => {
        val old: Int = map1.getOrElse(word, 0)
        map1.update(word, old + cnt)
      }
    }
  }*/

  override def merge(other: AccumulatorV2[String, mutable.Map[String, Int]]): Unit = {
    val map2: mutable.Map[String, Int] = other.value
    map2.foldLeft(map1){
      (map,kv)=>{
        val word: String = kv._1
        val cnt: Int = kv._2
        val old: Int = map.getOrElse(word, 0)
        map.update(word,old+cnt)
        map
      }
    }
  }

  override def value: mutable.Map[String, Int] = {
    map1
  }
}
